{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Data Preprocess"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "from utils import get_datasets\n",
    "from torchvision import transforms\n",
    "\n",
    "transforms = transforms.Compose([\n",
    "    transforms.Grayscale(num_output_channels=1),\n",
    "    transforms.ToTensor(),\n",
    "    # transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n",
    "])\n",
    "\n",
    "DATA_PATH = './data/'\n",
    "batch_size = 100\n",
    "\n",
    "train_dataset, test_dataset = get_datasets(DATA_PATH, transforms)\n",
    "train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
    "test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Network Design"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "\n",
    "class LeNet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Conv2d(1, 6, kernel_size=5, padding=2), nn.Sigmoid(),\n",
    "            nn.AvgPool2d(kernel_size=4, stride=3),\n",
    "            nn.Conv2d(6, 16, kernel_size=5), nn.Sigmoid(),\n",
    "            nn.AvgPool2d(kernel_size=4, stride=3),\n",
    "            nn.Flatten(), nn.Sigmoid(),\n",
    "        )\n",
    "        self.fc1 = nn.Sequential(\n",
    "            nn.Linear(1600, 120), nn.Sigmoid(),\n",
    "            nn.Linear(120, 26), nn.Softmax(dim=1)\n",
    "        )\n",
    "        self.fc2 = nn.Sequential(\n",
    "            nn.Linear(1600, 120), nn.Sigmoid(),\n",
    "            nn.Linear(120, 26), nn.Softmax(dim=1)\n",
    "        )\n",
    "        self.fc3 = nn.Sequential(\n",
    "            nn.Linear(1600, 120), nn.Sigmoid(),\n",
    "            nn.Linear(120, 26), nn.Softmax(dim=1)\n",
    "        )\n",
    "        self.fc4 = nn.Sequential(\n",
    "            nn.Linear(1600, 120), nn.Sigmoid(),\n",
    "            nn.Linear(120, 26), nn.Softmax(dim=1)\n",
    "        )\n",
    "        self.fc5 = nn.Sequential(\n",
    "            nn.Linear(1600, 120), nn.Sigmoid(),\n",
    "            nn.Linear(120, 26), nn.Softmax(dim=1)\n",
    "        )\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.net(x)\n",
    "        out1 = self.fc1(x)\n",
    "        out2 = self.fc2(x)\n",
    "        out3 = self.fc3(x)\n",
    "        out4 = self.fc4(x)\n",
    "        out5 = self.fc5(x)\n",
    "\n",
    "        out1 = out1.unsqueeze(1)\n",
    "        out2 = out2.unsqueeze(1)\n",
    "        out3 = out3.unsqueeze(1)\n",
    "        out4 = out4.unsqueeze(1)\n",
    "        out5 = out5.unsqueeze(1)\n",
    "\n",
    "        # (N, 5, 26)\n",
    "        out = torch.cat((out1, out2, out3, out4, out5), dim=1)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "========================================================================================================================\n",
       "Layer (type (var_name))                  Input Shape          Output Shape         Param #              Trainable\n",
       "========================================================================================================================\n",
       "LeNet (LeNet)                            [100, 1, 110, 110]   [100, 5, 26]         --                   True\n",
       "├─Sequential (net)                       [100, 1, 110, 110]   [100, 1600]          --                   True\n",
       "│    └─Conv2d (0)                        [100, 1, 110, 110]   [100, 6, 110, 110]   156                  True\n",
       "│    └─Sigmoid (1)                       [100, 6, 110, 110]   [100, 6, 110, 110]   --                   --\n",
       "│    └─AvgPool2d (2)                     [100, 6, 110, 110]   [100, 6, 36, 36]     --                   --\n",
       "│    └─Conv2d (3)                        [100, 6, 36, 36]     [100, 16, 32, 32]    2,416                True\n",
       "│    └─Sigmoid (4)                       [100, 16, 32, 32]    [100, 16, 32, 32]    --                   --\n",
       "│    └─AvgPool2d (5)                     [100, 16, 32, 32]    [100, 16, 10, 10]    --                   --\n",
       "│    └─Flatten (6)                       [100, 16, 10, 10]    [100, 1600]          --                   --\n",
       "│    └─Sigmoid (7)                       [100, 1600]          [100, 1600]          --                   --\n",
       "├─Sequential (fc1)                       [100, 1600]          [100, 26]            --                   True\n",
       "│    └─Linear (0)                        [100, 1600]          [100, 120]           192,120              True\n",
       "│    └─Sigmoid (1)                       [100, 120]           [100, 120]           --                   --\n",
       "│    └─Linear (2)                        [100, 120]           [100, 26]            3,146                True\n",
       "│    └─Softmax (3)                       [100, 26]            [100, 26]            --                   --\n",
       "├─Sequential (fc2)                       [100, 1600]          [100, 26]            --                   True\n",
       "│    └─Linear (0)                        [100, 1600]          [100, 120]           192,120              True\n",
       "│    └─Sigmoid (1)                       [100, 120]           [100, 120]           --                   --\n",
       "│    └─Linear (2)                        [100, 120]           [100, 26]            3,146                True\n",
       "│    └─Softmax (3)                       [100, 26]            [100, 26]            --                   --\n",
       "├─Sequential (fc3)                       [100, 1600]          [100, 26]            --                   True\n",
       "│    └─Linear (0)                        [100, 1600]          [100, 120]           192,120              True\n",
       "│    └─Sigmoid (1)                       [100, 120]           [100, 120]           --                   --\n",
       "│    └─Linear (2)                        [100, 120]           [100, 26]            3,146                True\n",
       "│    └─Softmax (3)                       [100, 26]            [100, 26]            --                   --\n",
       "├─Sequential (fc4)                       [100, 1600]          [100, 26]            --                   True\n",
       "│    └─Linear (0)                        [100, 1600]          [100, 120]           192,120              True\n",
       "│    └─Sigmoid (1)                       [100, 120]           [100, 120]           --                   --\n",
       "│    └─Linear (2)                        [100, 120]           [100, 26]            3,146                True\n",
       "│    └─Softmax (3)                       [100, 26]            [100, 26]            --                   --\n",
       "├─Sequential (fc5)                       [100, 1600]          [100, 26]            --                   True\n",
       "│    └─Linear (0)                        [100, 1600]          [100, 120]           192,120              True\n",
       "│    └─Sigmoid (1)                       [100, 120]           [100, 120]           --                   --\n",
       "│    └─Linear (2)                        [100, 120]           [100, 26]            3,146                True\n",
       "│    └─Softmax (3)                       [100, 26]            [100, 26]            --                   --\n",
       "========================================================================================================================\n",
       "Total params: 978,902\n",
       "Trainable params: 978,902\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (Units.MEGABYTES): 533.79\n",
       "========================================================================================================================\n",
       "Input size (MB): 4.84\n",
       "Forward/backward pass size (MB): 71.77\n",
       "Params size (MB): 3.92\n",
       "Estimated Total Size (MB): 80.53\n",
       "========================================================================================================================"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from torchinfo import summary\n",
    "\n",
    "model_lenet = LeNet()\n",
    "summary(model=model_lenet, input_size=(100, 1, 110, 110), col_width=20,\n",
    "                  col_names=['input_size', 'output_size', 'num_params', 'trainable'], row_settings=['var_names'], verbose=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Using wandb-core as the SDK backend.  Please refer to https://wandb.me/wandb-core for more information.\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mgdab\u001b[0m (\u001b[33mgdab-sjtu-hpc-center\u001b[0m). Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.19.1"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>d:\\Universe\\Earth\\24-7-10-SJTU-Appointment\\CaptchaRec\\wandb\\run-20241226_111600-533sxtm0</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href='https://wandb.ai/gdab-sjtu-hpc-center/SJTUCaptcha-Lenet/runs/533sxtm0' target=\"_blank\">splendid-universe-15</a></strong> to <a href='https://wandb.ai/gdab-sjtu-hpc-center/SJTUCaptcha-Lenet' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/developer-guide' target=\"_blank\">docs</a>)<br>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View project at <a href='https://wandb.ai/gdab-sjtu-hpc-center/SJTUCaptcha-Lenet' target=\"_blank\">https://wandb.ai/gdab-sjtu-hpc-center/SJTUCaptcha-Lenet</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run at <a href='https://wandb.ai/gdab-sjtu-hpc-center/SJTUCaptcha-Lenet/runs/533sxtm0' target=\"_blank\">https://wandb.ai/gdab-sjtu-hpc-center/SJTUCaptcha-Lenet/runs/533sxtm0</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<button onClick=\"this.nextSibling.style.display='block';this.style.display='none';\">Display W&B run</button><iframe src='https://wandb.ai/gdab-sjtu-hpc-center/SJTUCaptcha-Lenet/runs/533sxtm0?jupyter=true' style='border:none;width:100%;height:420px;display:none;'></iframe>"
      ],
      "text/plain": [
       "<wandb.sdk.wandb_run.Run at 0x2125031faa0>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "epochs = 100\n",
    "lr = 0.01\n",
    "\n",
    "# use wandb for log\n",
    "import wandb\n",
    "wandb.init(\n",
    "    project='SJTUCaptcha-Lenet',\n",
    "\n",
    "    config = {\n",
    "    'learning_rate': lr,\n",
    "    'architecture': 'LeNet',\n",
    "    'epochs': epochs,\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "training on cpu\n",
      "Epoch 1, Loss: 14.7243, Accuracy: 0.00%\n",
      "Epoch 2, Loss: 14.7211, Accuracy: 0.00%\n",
      "Epoch 3, Loss: 14.7189, Accuracy: 0.00%\n",
      "Epoch 4, Loss: 14.7174, Accuracy: 0.00%\n",
      "Epoch 5, Loss: 14.7162, Accuracy: 0.00%\n",
      "Epoch 6, Loss: 14.7153, Accuracy: 0.00%\n",
      "Epoch 7, Loss: 14.7146, Accuracy: 0.00%\n",
      "Epoch 8, Loss: 14.7139, Accuracy: 0.00%\n",
      "Epoch 9, Loss: 14.7134, Accuracy: 0.00%\n",
      "Epoch 10, Loss: 14.7129, Accuracy: 0.00%\n",
      "Epoch 11, Loss: 14.7126, Accuracy: 0.00%\n",
      "Epoch 12, Loss: 14.7122, Accuracy: 0.00%\n",
      "Epoch 13, Loss: 14.7120, Accuracy: 0.00%\n",
      "Epoch 14, Loss: 14.7118, Accuracy: 0.00%\n",
      "Epoch 15, Loss: 14.7116, Accuracy: 0.00%\n",
      "Epoch 16, Loss: 14.7114, Accuracy: 0.00%\n",
      "Epoch 17, Loss: 14.7112, Accuracy: 0.00%\n",
      "Epoch 18, Loss: 14.7111, Accuracy: 0.00%\n",
      "Epoch 19, Loss: 14.7109, Accuracy: 0.00%\n",
      "Epoch 20, Loss: 14.7108, Accuracy: 0.00%\n",
      "Epoch 21, Loss: 14.7107, Accuracy: 0.00%\n",
      "Epoch 22, Loss: 14.7105, Accuracy: 0.00%\n",
      "Epoch 23, Loss: 14.7104, Accuracy: 0.00%\n",
      "Epoch 24, Loss: 14.7104, Accuracy: 0.00%\n",
      "Epoch 25, Loss: 14.7102, Accuracy: 0.00%\n",
      "Epoch 26, Loss: 14.7101, Accuracy: 0.00%\n",
      "Epoch 27, Loss: 14.7100, Accuracy: 0.00%\n",
      "Epoch 28, Loss: 14.7100, Accuracy: 0.00%\n",
      "Epoch 29, Loss: 14.7099, Accuracy: 0.00%\n",
      "Epoch 30, Loss: 14.7098, Accuracy: 0.00%\n",
      "Epoch 31, Loss: 14.7097, Accuracy: 0.00%\n",
      "Epoch 32, Loss: 14.7096, Accuracy: 0.00%\n",
      "Epoch 33, Loss: 14.7095, Accuracy: 0.00%\n",
      "Epoch 34, Loss: 14.7094, Accuracy: 0.00%\n",
      "Epoch 35, Loss: 14.7093, Accuracy: 0.00%\n",
      "Epoch 36, Loss: 14.7093, Accuracy: 0.00%\n",
      "Epoch 37, Loss: 14.7092, Accuracy: 0.00%\n",
      "Epoch 38, Loss: 14.7091, Accuracy: 0.00%\n",
      "Epoch 39, Loss: 14.7089, Accuracy: 0.00%\n",
      "Epoch 40, Loss: 14.7089, Accuracy: 0.00%\n",
      "Epoch 41, Loss: 14.7087, Accuracy: 0.00%\n",
      "Epoch 42, Loss: 14.7087, Accuracy: 0.00%\n",
      "Epoch 43, Loss: 14.7087, Accuracy: 0.00%\n",
      "Epoch 44, Loss: 14.7086, Accuracy: 0.00%\n",
      "Epoch 45, Loss: 14.7086, Accuracy: 0.00%\n",
      "Epoch 46, Loss: 14.7085, Accuracy: 0.00%\n",
      "Epoch 47, Loss: 14.7085, Accuracy: 0.00%\n",
      "Epoch 48, Loss: 14.7084, Accuracy: 0.00%\n",
      "Epoch 49, Loss: 14.7084, Accuracy: 0.00%\n",
      "Epoch 50, Loss: 14.7083, Accuracy: 0.00%\n",
      "Epoch 51, Loss: 14.7082, Accuracy: 0.00%\n",
      "Epoch 52, Loss: 14.7082, Accuracy: 0.00%\n",
      "Epoch 53, Loss: 14.7081, Accuracy: 0.00%\n",
      "Epoch 54, Loss: 14.7080, Accuracy: 0.00%\n",
      "Epoch 55, Loss: 14.7080, Accuracy: 0.00%\n",
      "Epoch 56, Loss: 14.7079, Accuracy: 0.00%\n",
      "Epoch 57, Loss: 14.7078, Accuracy: 0.00%\n",
      "Epoch 58, Loss: 14.7079, Accuracy: 0.00%\n",
      "Epoch 59, Loss: 14.7078, Accuracy: 0.00%\n",
      "Epoch 60, Loss: 14.7077, Accuracy: 0.00%\n",
      "Epoch 61, Loss: 14.7078, Accuracy: 0.00%\n",
      "Epoch 62, Loss: 14.7077, Accuracy: 0.00%\n",
      "Epoch 63, Loss: 14.7077, Accuracy: 0.00%\n",
      "Epoch 64, Loss: 14.7077, Accuracy: 0.00%\n",
      "Epoch 65, Loss: 14.7077, Accuracy: 0.00%\n",
      "Epoch 66, Loss: 14.7077, Accuracy: 0.00%\n",
      "Epoch 67, Loss: 14.7077, Accuracy: 0.00%\n",
      "Epoch 68, Loss: 14.7077, Accuracy: 0.00%\n",
      "Epoch 69, Loss: 14.7076, Accuracy: 0.00%\n",
      "Epoch 70, Loss: 14.7076, Accuracy: 0.00%\n",
      "Epoch 71, Loss: 14.7076, Accuracy: 0.00%\n",
      "Epoch 72, Loss: 14.7075, Accuracy: 0.00%\n",
      "Epoch 73, Loss: 14.7076, Accuracy: 0.00%\n",
      "Epoch 74, Loss: 14.7076, Accuracy: 0.00%\n",
      "Epoch 75, Loss: 14.7075, Accuracy: 0.00%\n",
      "Epoch 76, Loss: 14.7075, Accuracy: 0.00%\n",
      "Epoch 77, Loss: 14.7075, Accuracy: 0.00%\n",
      "Epoch 78, Loss: 14.7075, Accuracy: 0.00%\n",
      "Epoch 79, Loss: 14.7075, Accuracy: 0.00%\n",
      "Epoch 80, Loss: 14.7075, Accuracy: 0.00%\n",
      "Epoch 81, Loss: 14.7075, Accuracy: 0.00%\n",
      "Epoch 82, Loss: 14.7075, Accuracy: 0.00%\n",
      "Epoch 83, Loss: 14.7075, Accuracy: 0.00%\n",
      "Epoch 84, Loss: 14.7075, Accuracy: 0.00%\n",
      "Epoch 85, Loss: 14.7074, Accuracy: 0.00%\n",
      "Epoch 86, Loss: 14.7074, Accuracy: 0.00%\n",
      "Epoch 87, Loss: 14.7074, Accuracy: 0.00%\n",
      "Epoch 88, Loss: 14.7074, Accuracy: 0.00%\n",
      "Epoch 89, Loss: 14.7074, Accuracy: 0.00%\n",
      "Epoch 90, Loss: 14.7074, Accuracy: 0.00%\n",
      "Epoch 91, Loss: 14.7073, Accuracy: 0.00%\n",
      "Epoch 92, Loss: 14.7074, Accuracy: 0.00%\n",
      "Epoch 93, Loss: 14.7073, Accuracy: 0.00%\n",
      "Epoch 94, Loss: 14.7073, Accuracy: 0.00%\n",
      "Epoch 95, Loss: 14.7074, Accuracy: 0.00%\n",
      "Epoch 96, Loss: 14.7073, Accuracy: 0.00%\n",
      "Epoch 97, Loss: 14.7073, Accuracy: 0.00%\n",
      "Epoch 98, Loss: 14.7073, Accuracy: 0.00%\n",
      "Epoch 99, Loss: 14.7073, Accuracy: 0.00%\n",
      "Epoch 100, Loss: 14.7073, Accuracy: 0.00%\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from utils import onehot_to_label\n",
    "\n",
    "def train_lenet(model, train_iter, test_iter, num_epochs, lr, device):\n",
    "    \"\"\"\n",
    "    Train Lemodel Model\n",
    "    - train_iter: training data loader\n",
    "    - test_iter: test data loader\n",
    "    - num_epochs: number of epochs\n",
    "    - lr: learning rate\n",
    "    - device: cpu or gpu\n",
    "    \"\"\"\n",
    "    def init_weights(m):\n",
    "        if type(m) == nn.Linear or type(m) == nn.Conv2d:\n",
    "            nn.init.xavier_uniform_(m.weight)\n",
    "    model.apply(init_weights)\n",
    "    print('training on', device)\n",
    "    model.to(device)\n",
    "\n",
    "    optimizer = torch.optim.SGD(model.parameters(), lr=lr)\n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "    # Training Loop\n",
    "    for epoch in range(num_epochs):\n",
    "        model.train()\n",
    "        train_loss, correct, total = 0, 0, 0\n",
    "        for X, y in train_iter:\n",
    "            optimizer.zero_grad()\n",
    "            X, y = X.to(device), y.to(device)\n",
    "            y_hat = model(X)\n",
    "            loss = []\n",
    "            for index in range(y_hat.size(1)):\n",
    "                loss.append(criterion(y_hat[:, index], y[:, index]))\n",
    "            loss = torch.stack(loss).sum()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "            # get predicted y\n",
    "            threshold = 0.5\n",
    "            predicted = torch.zeros_like(y_hat)\n",
    "            for i in range(y_hat.size(0)):\n",
    "                for j in range(y_hat.size(1)):\n",
    "                    _, max_idx = torch.max(y_hat[i, j], dim=0)\n",
    "                    if j == y_hat.size(1) - 1 and y_hat[i, j, max_idx] < threshold:\n",
    "                        predicted[i, j] = torch.zeros_like(y_hat[i, j])\n",
    "                    else:\n",
    "                        predicted[i, j, max_idx] = 1\n",
    "\n",
    "            # for i in range(y.size(0)):\n",
    "            #     print('real y:', onehot_to_label(y[i]))\n",
    "            #     print('predict y:', onehot_to_label(predicted[i]))\n",
    "\n",
    "            # for i in range(y.size(0)):\n",
    "            #     if (predicted[i] == y[i]).all():\n",
    "            #         print('real y:', onehot_to_label(y[i]))\n",
    "            #         print('predict y:', onehot_to_label(predicted[i]))\n",
    "   \n",
    "            train_loss += loss.item()\n",
    "            total += y.size(0)\n",
    "            for i in range(y.size(0)):\n",
    "                if (predicted[i] == y[i]).all():\n",
    "                    correct += 1\n",
    "            # print(f'num_correct: {correct}, total: {total}')\n",
    "\n",
    "            wandb.log({'train_loss': loss.item(), 'train_accuracy': correct / total})\n",
    "\n",
    "        print(f'Epoch {epoch+1}, Loss: {train_loss/len(train_iter):.4f}, Accuracy: {100 * correct / total:.2f}%')\n",
    "        \n",
    "    # # Validation Loop\n",
    "    # model.eval()\n",
    "    # test_loss, correct, total = 0, 0, 0\n",
    "    # with torch.no_grad():\n",
    "    #     for X, y in test_iter:\n",
    "    #         X, y = X.to(device), y.to(device)\n",
    "    #         y_hat = model(X)\n",
    "    #         loss = criterion(y_hat, y)\n",
    "    #         test_loss += loss.item()\n",
    "    #         _, predicted = torch.max(y_hat.data, 1)\n",
    "    #         total += y.size(0)\n",
    "    #         correct += (predicted == y).sum().item()\n",
    "\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "train_lenet(model_lenet, train_loader, test_loader, epochs, lr, device)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Save the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saving the model: models\\lenet5_v1_mnist.pth\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\1\\AppData\\Local\\Temp\\ipykernel_21420\\2021183831.py:15: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  model_lenet5_v1_mnist_loaded.load_state_dict(torch.load(MODEL_SAVE_PATH))\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "MODEL_PATH = Path(\"models\")\n",
    "MODEL_PATH.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "MODEL_NAME = \"lenet5_v1_mnist.pth\"\n",
    "MODEL_SAVE_PATH = MODEL_PATH / MODEL_NAME\n",
    "\n",
    "# Saving the model\n",
    "print(f\"Saving the model: {MODEL_SAVE_PATH}\")\n",
    "torch.save(obj=model_lenet.state_dict(), f=MODEL_SAVE_PATH)\n",
    "\n",
    "# Loading the saved model\n",
    "model_lenet5_v1_mnist_loaded = LeNet()\n",
    "model_lenet5_v1_mnist_loaded.load_state_dict(torch.load(MODEL_SAVE_PATH))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Evaluation on the test dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'test_dataloader' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[8], line 7\u001b[0m\n\u001b[0;32m      5\u001b[0m model_lenet5_v1_mnist_loaded\u001b[38;5;241m.\u001b[39meval()\n\u001b[0;32m      6\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m torch\u001b[38;5;241m.\u001b[39minference_mode():\n\u001b[1;32m----> 7\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m X, y \u001b[38;5;129;01min\u001b[39;00m \u001b[43mtest_dataloader\u001b[49m:\n\u001b[0;32m      8\u001b[0m         X, y \u001b[38;5;241m=\u001b[39m X\u001b[38;5;241m.\u001b[39mto(device), y\u001b[38;5;241m.\u001b[39mto(device)\n\u001b[0;32m      9\u001b[0m         y_pred \u001b[38;5;241m=\u001b[39m model_lenet5_v1_mnist_loaded(X)\n",
      "\u001b[1;31mNameError\u001b[0m: name 'test_dataloader' is not defined"
     ]
    }
   ],
   "source": [
    "test_loss, test_acc = 0, 0\n",
    "\n",
    "model_lenet5_v1_mnist_loaded.to(device)\n",
    "\n",
    "model_lenet5_v1_mnist_loaded.eval()\n",
    "with torch.inference_mode():\n",
    "    for X, y in test_dataloader:\n",
    "        X, y = X.to(device), y.to(device)\n",
    "        y_pred = model_lenet5_v1_mnist_loaded(X)\n",
    "        \n",
    "        test_loss += loss_fn(y_pred, y)\n",
    "        test_acc += accuracy(y_pred, y)\n",
    "        \n",
    "    test_loss /= len(test_dataloader)\n",
    "    test_acc /= len(test_dataloader)\n",
    "\n",
    "print(f\"Test loss: {test_loss: .5f}| Test acc: {test_acc: .5f}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
